# import paddlex as pdx
import numpy as np
import paddle
import pyodbc

from paddle.io import Dataset


class MyDataset(Dataset):  # 处理数据集
    """
    步骤一：继承paddle.io.Dataset
    """

    def __init__(self, dataname, dic_kw_path, dic_wk_path, batch_size=1):
        """
        步骤二：实现构造函数，定义数据读取方式，划分训练和测试数据集
        """
        super(MyDataset, self).__init__()

        self.dic_kw = self.Dic_read(dic_kw_path)
        self.dic_wk = self.Dic_read(dic_wk_path)
        connect = pyodbc.connect(
            r'DRIVER={ODBC Driver 17 for SQL Server};SERVER=10.23.66.8;DATABASE=语料库;UID=yyhsql;PWD=sql7850915')
        if connect:
            print("连接成功!")
        mycursor = connect.cursor()  # 创建一个游标对象,python里的sql语句都要通过cursor来执行
        sql = "select *  from " + dataname + ".语句表 where ID>2120"
        mycursor.execute(sql)
        self.rows = mycursor.fetchall()
        self.len = len(self.rows)

    def __getitem__(self, index):
        """
        步骤三：实现__getitem__方法，定义指定index时如何获取数据，并返回单条数据（训练数据，对应的标签）
        """
        input_word = self.rows[index][1]
        data, seq_len, maskd = self.get_tensor(str(input_word) , 100)
        label_word = self.rows[index][2]
        label, label_len, maskl = self.get_tensor(str(label_word) , 100)
        data=np.array(data, dtype="int64")
        seq_len=np.array(seq_len, dtype="int64")
        maskd=np.array(maskd,dtype="int64")
        label=np.array(label, dtype="int64")
        label_len=np.array(label_len,dtype="int64")
        maskl=np.array(maskl, dtype="int64")
        return data, seq_len, maskd,label,label,label_len, maskl


    def Dic_read(self, dic_name):
        dict = np.load(dic_name, allow_pickle='TRUE').item()
        return dict

    def get_tensor(self, word_str, maxlen):
        data_list = []
        mask = []
        # 0: '<PiD>', 1: '<puk>', 2: '<eos>', 3: '<SOS>',
        for char in word_str:

            # if char != " ":
            mask.append(1)
            if not char in self.dic_wk.keys():
                data_list.append(1)
            else:
                data_list.append(self.dic_wk[char])
            if len(data_list) == maxlen - 1:
                break
        data_list.append(2)
        mask.append(1)
        seq_len = len(data_list)
        while len(data_list) < maxlen:
            data_list.append(0)
            mask.append(0)

        return data_list, [seq_len], mask

    def get_word(self, ts):
        # ts是一个tensor变量，为batch_size,seq_len
        ts = ts.numpy()
        str_batch = []
        for i in range(ts.shape[0]):
            str_list = []
            for j in range(ts.shape[1]):
                str_list.append(self.dic_kw[ts[i][j]])
            str_batch.append(str_list)
        return str_batch

    def num_convert(self, num, n, w):
        num_list = []
        for i in range(w):
            num_list.append(int(num % n))
            num = num / n
        return num_list[::-1]

    def getrows(self):
        return self.rows

    def __len__(self):
        """
        步骤四：实现__len__方法，返回数据集总数目
        """
        return self.len
